#!/usr/bin/env python3
"""
Obstacle Detector - 障碍物检测器
基于Livox Mid360激光雷达点云数据进行障碍物检测
"""

import rclpy
from rclpy.node import Node
from rclpy.callback_groups import ReentrantCallbackGroup

from sensor_msgs.msg import PointCloud2
from geometry_msgs.msg import PoseStamped
from std_msgs.msg import Bool
from nav_msgs.msg import OccupancyGrid
from visualization_msgs.msg import MarkerArray, Marker

import numpy as np
import struct
from sensor_msgs_py import point_cloud2

class ObstacleDetector(Node):
    """
    障碍物检测器
    功能：
    1. 处理Livox Mid360点云数据
    2. 检测前方和周围障碍物
    3. 生成occupancy grid地图
    4. 发布障碍物检测结果
    """
    
    def __init__(self):
        super().__init__('obstacle_detector')
        
        # 检测参数
        self.declare_parameter('detection_range', 5.0)      # 检测范围(米)
        self.declare_parameter('safety_distance', 1.5)     # 安全距离(米)
        self.declare_parameter('min_obstacle_height', 0.2) # 最小障碍物高度(米)
        self.declare_parameter('max_obstacle_height', 3.0) # 最大障碍物高度(米)
        self.declare_parameter('grid_resolution', 0.1)     # 栅格地图分辨率(米)
        self.declare_parameter('grid_size', 100)           # 栅格地图大小
        
        self.detection_range = self.get_parameter('detection_range').value
        self.safety_distance = self.get_parameter('safety_distance').value
        self.min_obstacle_height = self.get_parameter('min_obstacle_height').value
        self.max_obstacle_height = self.get_parameter('max_obstacle_height').value
        self.grid_resolution = self.get_parameter('grid_resolution').value
        self.grid_size = self.get_parameter('grid_size').value
        
        # 状态变量
        self.current_pose = None
        self.obstacle_detected = False
        self.obstacle_grid = np.zeros((self.grid_size, self.grid_size), dtype=np.int8)
        
        # 回调组
        self.callback_group = ReentrantCallbackGroup()
        
        # 发布器
        self.obstacle_pub = self.create_publisher(
            Bool, '/obstacle_detected', 10)
        self.grid_pub = self.create_publisher(
            OccupancyGrid, '/obstacle_grid', 10)
        self.marker_pub = self.create_publisher(
            MarkerArray, '/obstacle_markers', 10)
            
        # 订阅器 - Livox Mid360点云话题
        self.pointcloud_sub = self.create_subscription(
            PointCloud2, '/livox/lidar', self.pointcloud_callback, 10,
            callback_group=self.callback_group)
        self.pose_sub = self.create_subscription(
            PoseStamped, '/current_pose', self.pose_callback, 10,
            callback_group=self.callback_group)
            
        # 创建定时器发布地图
        self.map_timer = self.create_timer(
            1.0, self.publish_occupancy_grid, callback_group=self.callback_group)
            
        self.get_logger().info("障碍物检测器已启动 - 监听Livox Mid360数据")
        
    def pose_callback(self, msg):
        """当前位置回调"""
        self.current_pose = msg
        
    def pointcloud_callback(self, msg):
        """点云数据回调"""
        if self.current_pose is None:
            return
            
        try:
            # 提取点云数据
            points = self.extract_points_from_pointcloud(msg)
            
            if len(points) == 0:
                return
                
            # 过滤点云（移除地面和过高的点）
            filtered_points = self.filter_pointcloud(points)
            
            # 障碍物检测
            obstacles_detected = self.detect_obstacles(filtered_points)
            
            # 更新占用栅格地图
            self.update_occupancy_grid(filtered_points)
            
            # 发布检测结果
            obstacle_msg = Bool()
            obstacle_msg.data = obstacles_detected
            self.obstacle_pub.publish(obstacle_msg)
            
            # 可视化障碍物
            self.visualize_obstacles(filtered_points)
            
            # 更新状态
            if obstacles_detected != self.obstacle_detected:
                self.obstacle_detected = obstacles_detected
                if obstacles_detected:
                    self.get_logger().warn("检测到障碍物！")
                else:
                    self.get_logger().info("障碍物消失")
                    
        except Exception as e:
            self.get_logger().error(f"点云处理错误: {str(e)}")
            
    def extract_points_from_pointcloud(self, pointcloud_msg):
        """从PointCloud2消息中提取点云数据"""
        points = []
        
        # 使用sensor_msgs_py库解析点云
        for point in point_cloud2.read_points(pointcloud_msg, 
                                             field_names=("x", "y", "z"), 
                                             skip_nans=True):
            points.append([point[0], point[1], point[2]])
            
        return np.array(points)
        
    def filter_pointcloud(self, points):
        """过滤点云数据 - 优化版本使用numpy向量化操作"""
        if len(points) == 0:
            return points
            
        # 使用numpy向量化操作提升性能
        points = np.array(points)
        
        # 距离过滤 - 向量化计算
        distances = np.sqrt(points[:, 0]**2 + points[:, 1]**2)
        distance_mask = (distances <= self.detection_range) & (distances >= 0.5)
        
        # 高度过滤 - 向量化操作
        height_mask = (points[:, 2] >= self.min_obstacle_height) & (points[:, 2] <= self.max_obstacle_height)
        
        # 组合过滤条件
        valid_mask = distance_mask & height_mask
        
        return points[valid_mask]
        
    def detect_obstacles(self, points):
        """检测前方是否有障碍物 - 优化版本"""
        if len(points) == 0:
            return False
            
        # 向量化计算距离和角度
        distances = np.sqrt(points[:, 0]**2 + points[:, 1]**2)
        angles = np.arctan2(points[:, 1], points[:, 0])
        
        # 前方扇形区域检测（±45度）
        front_mask = (np.abs(angles) < np.pi/4) & (distances < self.safety_distance)
        obstacle_count = np.sum(front_mask)
        
        # 如果检测到足够多的障碍点，认为有障碍物
        return obstacle_count > 10
        
    def update_occupancy_grid(self, points):
        """更新占用栅格地图 - 优化版本"""
        # 重置地图
        self.obstacle_grid.fill(-1)  # -1表示未知
        
        if len(points) == 0:
            return
            
        # 地图中心对应无人机当前位置
        center_x = self.grid_size // 2
        center_y = self.grid_size // 2
        
        # 向量化坐标转换
        grid_x = ((points[:, 0] / self.grid_resolution) + center_x).astype(int)
        grid_y = ((points[:, 1] / self.grid_resolution) + center_y).astype(int)
        
        # 过滤超出范围的点
        valid_indices = ((grid_x >= 0) & (grid_x < self.grid_size) & 
                        (grid_y >= 0) & (grid_y < self.grid_size))
        
        valid_x = grid_x[valid_indices]
        valid_y = grid_y[valid_indices]
        
        # 批量更新网格
        self.obstacle_grid[valid_y, valid_x] = 100  # 100表示有障碍物
        
        # 膨胀操作（可选，用于安全边界）
        self.dilate_obstacles()
        
    def dilate_obstacles(self):
        """膨胀障碍物 - 优化版本使用scipy或手动优化"""
        import scipy.ndimage
        
        # 创建膨胀核（圆形）
        kernel_size = max(1, int(self.safety_distance / self.grid_resolution))
        kernel = np.ones((kernel_size*2+1, kernel_size*2+1), dtype=bool)
        
        # 创建圆形膨胀核
        center = kernel_size
        y, x = np.ogrid[:kernel_size*2+1, :kernel_size*2+1]
        mask = (x - center)**2 + (y - center)**2 <= kernel_size**2
        kernel = kernel & mask
        
        # 找到所有障碍物位置
        obstacle_mask = (self.obstacle_grid == 100)
        
        if np.any(obstacle_mask):
            # 使用scipy进行高效膨胀
            try:
                dilated = scipy.ndimage.binary_dilation(obstacle_mask, kernel)
                self.obstacle_grid[dilated] = 100
            except ImportError:
                # 备用方案：手动膨胀
                self._manual_dilation(obstacle_mask, kernel_size)
                
    def _manual_dilation(self, obstacle_mask, kernel_size):
        """手动膨胀操作（当scipy不可用时）"""
        obstacle_positions = np.argwhere(obstacle_mask)
        
        for pos in obstacle_positions:
            y, x = pos
            y_min = max(0, y - kernel_size)
            y_max = min(self.grid_size, y + kernel_size + 1)
            x_min = max(0, x - kernel_size)
            x_max = min(self.grid_size, x + kernel_size + 1)
            
            # 在安全半径内标记为障碍物
            yy, xx = np.ogrid[y_min:y_max, x_min:x_max]
            circle_mask = (xx - x)**2 + (yy - y)**2 <= kernel_size**2
            self.obstacle_grid[y_min:y_max, x_min:x_max][circle_mask] = 100
            
    def publish_occupancy_grid(self):
        """发布占用栅格地图"""
        if self.current_pose is None:
            return
            
        grid_msg = OccupancyGrid()
        grid_msg.header.frame_id = "map"
        grid_msg.header.stamp = self.get_clock().now().to_msg()
        
        # 地图信息
        grid_msg.info.resolution = self.grid_resolution
        grid_msg.info.width = self.grid_size
        grid_msg.info.height = self.grid_size
        
        # 地图原点（无人机位置为中心）
        grid_msg.info.origin.position.x = (
            self.current_pose.pose.position.x - 
            (self.grid_size * self.grid_resolution) / 2
        )
        grid_msg.info.origin.position.y = (
            self.current_pose.pose.position.y - 
            (self.grid_size * self.grid_resolution) / 2
        )
        grid_msg.info.origin.position.z = 0.0
        grid_msg.info.origin.orientation.w = 1.0
        
        # 地图数据
        grid_msg.data = self.obstacle_grid.flatten().tolist()
        
        self.grid_pub.publish(grid_msg)
        
    def visualize_obstacles(self, points):
        """可视化障碍物 - 优化版本"""
        if len(points) == 0:
            # 清空标记
            empty_markers = MarkerArray()
            empty_marker = Marker()
            empty_marker.action = Marker.DELETEALL
            empty_markers.markers.append(empty_marker)
            self.marker_pub.publish(empty_markers)
            return
            
        marker_array = MarkerArray()
        
        # 下采样点云以减少可视化负载
        if len(points) > 1000:
            # 随机下采样到1000个点
            indices = np.random.choice(len(points), 1000, replace=False)
            points = points[indices]
        
        # 创建点云标记
        marker = Marker()
        marker.header.frame_id = "base_link"
        marker.header.stamp = self.get_clock().now().to_msg()
        marker.ns = "obstacles"
        marker.id = 0
        marker.type = Marker.POINTS
        marker.action = Marker.ADD
        marker.pose.orientation.w = 1.0
        
        # 设置点的大小和颜色
        marker.scale.x = 0.05
        marker.scale.y = 0.05
        marker.color.r = 1.0
        marker.color.g = 0.0
        marker.color.b = 0.0
        marker.color.a = 0.8
        
        # 批量添加点
        from geometry_msgs.msg import Point
        from std_msgs.msg import ColorRGBA
        
        for point in points:
            p = Point()
            p.x, p.y, p.z = float(point[0]), float(point[1]), float(point[2])
            marker.points.append(p)
            
            # 根据距离设置颜色
            distance = np.sqrt(point[0]**2 + point[1]**2)
            color = ColorRGBA()
            if distance < self.safety_distance:
                color.r, color.g, color.b, color.a = 1.0, 0.0, 0.0, 1.0  # 红色
            else:
                color.r, color.g, color.b, color.a = 1.0, 1.0, 0.0, 0.8  # 黄色
            marker.colors.append(color)
        
        marker_array.markers.append(marker)
        
        # 添加安全区域可视化
        safety_marker = Marker()
        safety_marker.header.frame_id = "base_link"
        safety_marker.header.stamp = self.get_clock().now().to_msg()
        safety_marker.ns = "safety_zone"
        safety_marker.id = 1
        safety_marker.type = Marker.CYLINDER
        safety_marker.action = Marker.ADD
        
        safety_marker.pose.position.x = 0.0
        safety_marker.pose.position.y = 0.0
        safety_marker.pose.position.z = 1.0
        safety_marker.pose.orientation.w = 1.0
        
        safety_marker.scale.x = self.safety_distance * 2
        safety_marker.scale.y = self.safety_distance * 2
        safety_marker.scale.z = 0.1
        
        safety_marker.color.r = 0.0
        safety_marker.color.g = 1.0
        safety_marker.color.b = 0.0
        safety_marker.color.a = 0.2
        
        marker_array.markers.append(safety_marker)
        
        self.marker_pub.publish(marker_array)

def main(args=None):
    rclpy.init(args=args)
    
    obstacle_detector = ObstacleDetector()
    
    try:
        rclpy.spin(obstacle_detector)
    except KeyboardInterrupt:
        pass
    finally:
        obstacle_detector.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main() 